package com.wgz.ai.langchain4jai.memory;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;

import static dev.langchain4j.data.message.UserMessage.userMessage;
/**
 * @Title: WGZ
 * @Package com.wgz.ai.langchain4jai.memory
 * @Date 2025/3/5 23:49
 * @description: 带有聊天记忆的对话链
 */
public class ChatMemoryExamples {
    private final static String API_KEY = "sk-svcacct-J-zoI616zoIN8ghM6amdG_SbFNc3vADd5yRa-53EQkr-lPN0tjD8QgWJfAmSuT3BlbkFJxBhbhopzngM8ZUB8KwByRFPnJ5bDUT1eq2Kop0OlggeqfzdsNSF45DQFZ94mgA";

    public static void main(String[] args) {
        // 创建一个聊天记忆对象，最大记忆 300 个 token，使用 OpenAI 的 tokenizer 计算 token 数量
        ChatMemory chatMemory = TokenWindowChatMemory.withMaxTokens(300, new OpenAiTokenizer("gpt-3.5-turbo"));

        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(API_KEY)
                .modelName("gpt-3.5-turbo")
                .build();

        // You have full control over the chat memory.
        // You can decide if you want to add a particular message to the memory
        // (e.g. you might not want to store few-shot examples to save on tokens).
        // You can process/modify the message before saving if required.

        // 添加用户消息 "Hello, my name is Klaus" 到聊天记忆
        chatMemory.add(userMessage("Hello, my name is Klaus"));
        AiMessage answer = model.chat(chatMemory.messages()).aiMessage();
        System.out.println(answer.text()); // Hello Klaus! How can I assist you today?
//        将 AI 的回复添加到聊天记忆，保持对话的连贯性
        chatMemory.add(answer);

        chatMemory.add(userMessage("What is my name?"));
        AiMessage answerWithName = model.chat(chatMemory.messages()).aiMessage();
        System.out.println(answerWithName.text()); // Your name is Klaus.
        chatMemory.add(answerWithName);
    }
}
